import os
import re
import time
from src.ui.i18n import _
from PySide6.QtCore import Signal
from src.ui.base_worker import BaseWorker
from src.utils.srt_parser import SRTParser
from src.translation.openai_translator import OpenAITranslator, TranslationError
from src.translation.microsoft_translator import MicrosoftTranslator
from src.utils.logger import get_logger



def translate_srt_core(worker, srt_path, translator_name, source_lang, target_lang, model_config=None, video_type="常规", translation_style="常规", no_think=False, progress_callback=None, stream_callback=None, batch_start_callback=None, batch_done_callback=None, all_done_callback=None, timing_callback=None, is_stopped_func=None, time_tick_callback=None, logger=None):
    """Core logic for subtitle translation."""
    total_start_time = time.time()
    logger = logger or get_logger(__name__)
    is_stopped = is_stopped_func or (lambda: False)

    last_tick_time = 0
    def tick_time():
        nonlocal last_tick_time
        now = time.time()
        if time_tick_callback and (now - last_tick_time > 0.2):
            elapsed_time = now - total_start_time
            time_tick_callback(elapsed_time, 0)
            last_tick_time = now
    
    def progress(msg):
        if progress_callback: progress_callback(msg)
        logger.info(msg)

    def record_time(stage, elapsed_time):
        logger.debug(_("⏱️  [Timing] {stage} took: {elapsed_time:.2f} s").format(stage=stage, elapsed_time=elapsed_time))
        if timing_callback:
            timing_callback(stage, elapsed_time)

    if is_stopped(): return None
    progress(_("Starting to parse SRT file for translation..."))
    stage_start_time = time.time()
    parser = SRTParser(srt_path)
    parser.parse()
    original_texts = parser.get_texts()
    record_time(_("Parse SRT file"), time.time() - stage_start_time)
    progress(_("File parsed successfully, found {} subtitles.").format(len(original_texts)))

    if is_stopped(): return None
    progress(_("Initializing translation service..."))
    stage_start_time = time.time()
    if translator_name == OpenAITranslator.get_name():
        translator = OpenAITranslator()
    elif translator_name == MicrosoftTranslator.get_name():
        translator = MicrosoftTranslator()
    else:
        raise NotImplementedError(_("Translation service '{}' is not supported yet.").format(translator_name))
    record_time(_("Initialize translator"), time.time() - stage_start_time)

    def internal_batch_callback(current_batch, total_batches, batch_size=0):
        if is_stopped():
            raise InterruptedError("Translation was cancelled by user.")
        if batch_done_callback:
             batch_done_callback(current_batch, total_batches) # Note: batch_size is not used here, but the signature is fixed
        tick_time()

    def internal_stream_callback(chunk):
        if stream_callback: stream_callback(chunk)
        tick_time()

    if is_stopped(): return None
    progress(_("Starting to translate text..."))
    stage_start_time = time.time()
    
    translate_args = {
        "texts": original_texts,
        "source_lang": source_lang,
        "target_lang": target_lang,
        "video_type": video_type,
        "translation_style": translation_style,
        "batch_start_callback": batch_start_callback,
        "batch_done_callback": internal_batch_callback,
        "stream_callback": internal_stream_callback,
        "is_cancelled": is_stopped
    }
    if translator_name == OpenAITranslator.get_name():
        if model_config is not None:
            translate_args["model_config"] = model_config
        translate_args["no_think"] = no_think

    translated_texts = translator.translate(**translate_args)
    record_time(_("Translate text"), time.time() - stage_start_time)
    
    if is_stopped(): return None

    progress(_("Translation finished, got {} translated lines.").format(len(translated_texts)))
    if len(original_texts) != len(translated_texts):
        progress(_("[Warning] Mismatch between original and translated text count! Original: {}, Translated: {}").format(len(original_texts), len(translated_texts)))
    
    if all_done_callback:
        all_done_callback(list(zip(original_texts, translated_texts)))
    
    progress(_("All batches translated, generating new SRT file..."))
    stage_start_time = time.time()
    base, ext = os.path.splitext(srt_path)

    # Define a mapping for simple codes to full locale names for consistency
    LOCALE_MAP = {
        "zh": "zh-CN", "en": "en-US", "ja": "ja-JP", "ko": "ko-KR",
        "fr": "fr-FR", "de": "de-DE", "ru": "ru-RU"
    }
    # Get the full locale for the target language
    target_locale = LOCALE_MAP.get(target_lang, target_lang)

    # Regex to find language codes like .en or .en-US at the end of the filename (before extension)
    lang_code_pattern = re.compile(r'\.([a-z]{2,3}(-[A-Z]{2})?)$')

    # Remove any existing language code from the base filename
    base = lang_code_pattern.sub('', base)

    output_path = f"{base}.{target_locale}{ext}"
    parser.write(output_path, translated_texts)
    record_time(_("Generate SRT file"), time.time() - stage_start_time)
    
    progress(_("Translation successful! Output file: {}").format(output_path))
    
    total_elapsed_time = time.time() - total_start_time
    record_time(_("Total time"), total_elapsed_time)
    
    return output_path

class InterruptedError(Exception):
    pass

class TranslationWorker(BaseWorker):
    time_updated = Signal(float, float)
    translation_chunk = Signal(str)
    translation_done = Signal(list)
    batch_started = Signal(int, int)
    batch_completed = Signal(int, int)
    timing_updated = Signal(str, float)
    task_finished = Signal(str, str)

    def __init__(self, srt_path, translator_name, source_lang, target_lang, model_config=None, video_type="常规", translation_style="常规", no_think=False, logger=None):
        super().__init__()
        self.output_path = None
        self.srt_path = srt_path
        self.translator_name = translator_name
        self.source_lang = source_lang
        self.target_lang = target_lang
        self.model_config = model_config
        self.video_type = video_type
        self.translation_style = translation_style
        self.no_think = no_think
        self.logger = logger or get_logger(__name__)

    def core_run(self):
        try:
            self.output_path = translate_srt_core(
                self,
                self.srt_path,
                self.translator_name,
                self.source_lang,
                self.target_lang,
                model_config=self.model_config,
                video_type=self.video_type,
                translation_style=self.translation_style,
                no_think=self.no_think,
                progress_callback=self.progress.emit,
                stream_callback=self.translation_chunk.emit,
                batch_start_callback=self.batch_started.emit,
                batch_done_callback=self.batch_completed.emit,
                all_done_callback=self.translation_done.emit,
                timing_callback=self.timing_updated.emit,
                is_stopped_func=self.isInterruptionRequested,
                time_tick_callback=self.time_updated.emit,
                logger=self.logger
            )
            
            if self.isInterruptionRequested():
                self.progress.emit(_("Translation was manually stopped by the user."))
                return

            if self.output_path:
                self.task_finished.emit(_("Translation successful!"), self.output_path)
            else:
                if not self.isInterruptionRequested():
                    raise Exception(_("Translation core logic did not return an output path and did not report an error."))
        except TranslationError as e:
            self.error.emit(str(e))
            raise  # Re-raise the exception to be caught by the batch worker
        except Exception as e:
            self.error.emit(_("An unknown error occurred: {}").format(e))
            raise  # Re-raise the exception to be caught by the batch worker